//
// Created by yk120 on 2024/3/1.
//

#ifndef NANOFAISS_INDEXIVF_H
#define NANOFAISS_INDEXIVF_H

#include <nanofaiss/Clustering.h>
#include <nanofaiss/Index.h>
#include <nanofaiss/invlists/InvertedLists.h>
#include <nanofaiss/invlists/DirectMap.h>


namespace faiss {

/** Encapsulate a quantizer object for IndexIVF
 */
struct Level1Quantizer {

    /// quantizer maps vector to inverted list
    Index *quantizer = nullptr;

    /// number of inverted list
    size_t nlist = 0;

    /**
     * = 0: use the quantizer as index in a k-means training
     */
    uint8_t quantizer_train_alone = 0;

    bool own_fields = false; /// whether object owns quantizer

    /// parameter of cluster
    ClusteringParameters cp;

    Index *clustering_index = nullptr;

    /** Trains the quantizer and calls train residual to train sub-quantizer
     * @param n         size of of train set
     * @param x         vector x
     * @param metric    metric type
     */
    void train_q1(
            size_t n,
            const float *x,
            MetricType metric);



    Level1Quantizer(Index *quantizer, size_t nlist);

    Level1Quantizer();

    virtual ~Level1Quantizer();
};


struct IndexIVFInterface : Level1Quantizer {
    size_t nprobe = 1; /// number of probes at query time
    size_t max_codes = 0; /// max nb of codes to visit to do a query

    explicit IndexIVFInterface(Index *quantizer = nullptr, size_t nlist = 0)
                : Level1Quantizer(quantizer, nlist) {}


    virtual ~IndexIVFInterface() {}
};


struct IndexIVF : Index, IndexIVFInterface {
    /// Access to the actual data
    InvertedLists* invlists = nullptr;
    bool own_invlists = false;

    size_t code_size = 0; /// code size per vector in bytes

    /** Parallel mode determines how queries
     *
     */
    int parallel_mode = 0;
    const int PARALLEL_MODE_NO_HEAP_INIT = 1024;

    DirectMap direct_map;

    /// do the codes in invlists encode the vector relative to
    /// the centroids
    bool by_residual = true;

    /** The Inverted file takes a quantizer(an Index) on input
     *  which implements the function mapping a vector to a list
     */
    IndexIVF(Index *quantizer,
             size_t d,
             size_t nlist,
             size_t code_size,
             MetricType metric = MetricType::METRIC_L2);


    void reset() override;

    /// Trains the quantizer and calls train_encoder to train sub-quantizer
    void train(idx_t n, const float *x) override;


    /// Calls add_with_ids with NULL ids
    void add(idx_t n, const float *x) override;


    /// default implementation that calls encode_vectors
    void add_with_ids(idx_t n, const float *x, const idx_t *xids) override;

    /// redefined by subclass to indicate how many training vectors they need
    virtual idx_t train_encoder_num_vectors() const;


    /** TODO
     * @param n
     * @param x
     * @param xids
     * @param precomputed_idx
     */
    virtual void add_core(
            idx_t n,
            const float *x,
            const idx_t *xids,
            const idx_t *precomputed_idx);

    /** Train the encoder of the vectors
     * @param n
     * @param x
     * @param assign
     */
    virtual void train_encoder(idx_t n, const float *x, const idx_t *assign);



};


}

#endif // NANOFAISS_INDEXIVF_H
